# improved_classifier.py
import numpy as np
import tensorflow as tf
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
import os
from improved_preprocessing import extract_detailed_features
from training_visualization import plot_training_history

def load_and_retrain():
    # 加载原始序列数据
    sequences = ['swipe_down', 'swipe_up', 'grab']
    X_new = []
    y_new = []

    for label, seq_name in enumerate(sequences):
        seq_dir = f'sequence_data/{seq_name}'
        for file in os.listdir(seq_dir):
            if file.endswith('.npy'):
                sequence_data = np.load(f'{seq_dir}/{file}')
                # 使用改进的特征提取
                features = extract_detailed_features(sequence_data)
                X_new.append(features)
                y_new.append(label)

    X_new = np.array(X_new)
    y_new = np.array(y_new)

    print(f"新特征数据集形状: X={X_new.shape}, y={y_new.shape}")

    # 数据标准化
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X_new)

    # 划分数据集
    X_train, X_test, y_train, y_test = train_test_split(
        X_scaled, y_new, test_size=0.2, random_state=42, stratify=y_new
    )

    # 创建改进的模型
    def create_improved_model(input_dim=80):
        model = tf.keras.Sequential([
            tf.keras.layers.Dense(256, activation='relu', input_shape=(input_dim,)),
            tf.keras.layers.BatchNormalization(),
            tf.keras.layers.Dropout(0.4),

            tf.keras.layers.Dense(128, activation='relu'),
            tf.keras.layers.BatchNormalization(),
            tf.keras.layers.Dropout(0.3),

            tf.keras.layers.Dense(64, activation='relu'),
            tf.keras.layers.Dropout(0.2),

            tf.keras.layers.Dense(32, activation='relu'),
            tf.keras.layers.Dropout(0.1),

            tf.keras.layers.Dense(3, activation='softmax')
        ])

        model.compile(
            optimizer=tf.keras.optimizers.Adam(learning_rate=0.001),
            loss='sparse_categorical_crossentropy',
            metrics=['accuracy']
        )

        return model

    # 训练模型
    model = create_improved_model()

    callbacks = [
        tf.keras.callbacks.EarlyStopping(patience=20, restore_best_weights=True),
        tf.keras.callbacks.ReduceLROnPlateau(factor=0.5, patience=10)
    ]

    history = model.fit(
        X_train, y_train,
        epochs=150,
        batch_size=32,
        validation_split=0.2,
        callbacks=callbacks,
        verbose=1
    )

    # 评估模型
    test_loss, test_accuracy = model.evaluate(X_test, y_test, verbose=0)
    print(f"改进模型测试准确率: {test_accuracy:.4f}")

    # 保存新模型
    model.save('improved_gesture_classifier.keras')
    np.save('improved_scaler_params.npy', {'mean': scaler.mean_, 'scale': scaler.scale_})

    # 分析每个类别的准确率
    from sklearn.metrics import classification_report
    y_pred = model.predict(X_test)
    y_pred_classes = np.argmax(y_pred, axis=1)

    print("\n分类报告:")
    print(classification_report(y_test, y_pred_classes,
                                target_names=['向下滑动', '向上滑动', '抓取动作']))

    # 绘制训练历史图表
    plot_training_history(history, "改进模型", "improved_model_training_history.png")

    return model, history

if __name__ == "__main__":
    model, history = load_and_retrain()
